package api;

import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.utils.Utils;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class CustomParititioner implements Partitioner {
    private Map<String,Integer> counter=new HashMap<String, Integer>();
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        System.out.println("partition:"+key);
        List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);
        int numPartitions = partitions.size();
        if(key==null){
            System.out.println("keyBytes==null");
            return nextCounter(counter,topic)%numPartitions;
        }else{
            System.out.println("keyBytes!=null");
            return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;
        }
    }

    public void close() {

    }
    public Integer nextCounter(Map<String,Integer> counter,String topic){
        if(counter.containsKey(topic)){
            int c= counter.get(topic);
            if(c+1<0){
                c=0;
            }
            counter.put(topic,c+1);
            return c;
        }else{
            counter.put(topic,0);
            return 0;
        }
    }
    public void configure(Map<String, ?> configs) {

    }
}
